#!/usr/bin/env python3
"""
Individual PBPK Modeling Script (Simplified Version)

This script uses existing modules in the src folder to implement PBPK modeling with
different absorption and distribution modes:
- Absorption modes: first_order, adam
- Distribution modes: minimal_pbpk, full_pbpk

Author: BlackCat
Version: 1.0.0
"""

import numpy as np
import matplotlib.pyplot as plt
import sys
import os

# Add src path to system path
# sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../src'))
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
print(f"Project root path: {project_root}")
sys.path.insert(0, project_root)

from src.absorption.absorption_module import AbsorptionModel
from src.distribution.distribution_module import DistributionModel
from src.elimination.elimination_module import EliminationModel


class SimplePBPKSimulator:
    """Simplified Individual PBPK Simulator"""
    
    def __init__(self):
        """Initialize simulator"""
        self.default_params = self._get_default_parameters()
        self.time_points = np.linspace(0, 24, 48)  # 24小时，48个时间点
        
    def _get_default_parameters(self):
        """Get default parameters"""
        return {
            # Basic parameters
            'dose': 100.0,  # mg
            'body_weight': 70.0,  # kg
            
            # Absorption parameters
            'ka': 0.5,  # 吸收速率常数 (1/h)
            'fa': 0.8,  # 吸收分数
            'fg': 0.9,  # 肠道代谢分数
            'fh': 0.7,  # 肝脏代谢分数
            
            # Distribution parameters
            'Vp': 3.0,  # 血浆容积 (L)
            'Ve': 12.0,  # 细胞外液容积 (L)
            'Vt': 12.0,  # 组织容积 (L)
            'fu': 0.1,  # 游离分数
            'BP': 1.0,  # 血液-血浆浓度比
            
            # 组织分配系数
            'Kp_liver': 2.0,
            'Kp_kidney': 2.5,
            'Kp_brain': 0.5,
            'Kp_lung': 1.5,
            'Kp_heart': 1.8,
            'Kp_muscle': 1.2,
            'Kp_adipose': 0.8,
            
            # Elimination parameters
            'CL_hepatic': 0.5,  # 肝脏清除率 (L/h)
            'CL_renal': 0.3,  # 肾脏清除率 (L/h)
            'CL_biliary': 0.1,  # 胆汁清除率 (L/h)
            'CL_other': 0.05,  # 其他清除率 (L/h)
        }
    
    def run_first_order_absorption(self, params, time_points):
        """Run first-order absorption model"""
        model = AbsorptionModel(params)
        result = model.calculate_absorption(params, time_points)
        return {
            'plasma_concentration': result['plasma_concentration'],
            'cumulative_absorption': result['cumulative_absorption'],
            'absorption_rate': result['absorption_rate'],
            'bioavailability': result['bioavailability']
        }
    
    def run_adam_absorption(self, params, time_points):
        """Run ADAM absorption model"""
        # 简化的ADAM模型实现
        ka = params.get('ka', 0.5)
        fa = params.get('fa', 0.8)
        dose = params.get('dose', 100.0)
        
        # 简单的ADAM模型近似
        cumulative_absorption = dose * fa * (1 - np.exp(-ka * time_points))
        absorption_rate = dose * fa * ka * np.exp(-ka * time_points)
        
        # 简化的血浆浓度计算
        Vd = params.get('Vp', 3.0) + params.get('Vt', 12.0)
        plasma_concentration = absorption_rate / (ka * Vd) * (1 - np.exp(-ka * time_points))
        
        return {
            'plasma_concentration': plasma_concentration,
            'cumulative_absorption': cumulative_absorption,
            'absorption_rate': absorption_rate,
            'bioavailability': fa
        }
    
    def run_distribution(self, model_type, params, time_points, absorption_result):
        """Run distribution model"""
        absorption_input = {
            'dose': params['dose'],
            'cumulative': absorption_result['cumulative_absorption'],
            'plasma_concentration': absorption_result['plasma_concentration']
        }
        
        model = DistributionModel(params)
        result = model.calculate_distribution(params, time_points, absorption_input)
        
        if model_type == 'full_pbpk':
            # 计算脑部分布
            brain_result = model.calculate_brain_distribution(params, time_points, result['plasma_concentration'])
            return {
                'plasma_concentration': result['plasma_concentration'],
                'distribution_volume': result['distribution_volume'],
                'tissue_concentrations': result['tissue_concentrations'],
                'brain_concentration': brain_result['concentration']
            }
        else:
            return {
                'plasma_concentration': result['plasma_concentration'],
                'distribution_volume': result['distribution_volume'],
                'tissue_concentrations': result['tissue_concentrations']
            }
    
    def run_elimination(self, params, time_points, distribution_result):
        """Run elimination model"""
        model = EliminationModel(params)
        
        # 简化的消除计算
        plasma_conc = distribution_result['plasma_concentration']
        CL_total = params.get('CL_hepatic', 0.5) + params.get('CL_renal', 0.3) + params.get('CL_biliary', 0.1)
        Vd = distribution_result.get('distribution_volume', 15.0)
        ke = CL_total / Vd  # 消除速率常数
        
        # 计算消除后的浓度
        eliminated_conc = plasma_conc * np.exp(-ke * time_points)
        
        return {
            'plasma_concentration': eliminated_conc,
            'clearance_rates': {
                'total': CL_total,
                'hepatic': params.get('CL_hepatic', 0.5),
                'renal': params.get('CL_renal', 0.3),
                'biliary': params.get('CL_biliary', 0.1)
            }
        }
    
    def run_complete_pbpk(self, absorption_model, distribution_model, params=None):
        """Run complete PBPK simulation"""
        if params is None:
            params = self.default_params.copy()
        
        print(f"\nRunning simulation: absorption={absorption_model}, distribution={distribution_model}")
        
        # 1. 吸收阶段
        if absorption_model == 'first_order':
            absorption_result = self.run_first_order_absorption(params, self.time_points)
        else:  # adam
            absorption_result = self.run_adam_absorption(params, self.time_points)
        
        print(f"  Absorption completed: bioavailability={absorption_result['bioavailability']:.2%}")
        
        # 2. 分布阶段
        distribution_result = self.run_distribution(distribution_model, params, self.time_points, absorption_result)
        print(f"  Distribution completed: volume of distribution={distribution_result['distribution_volume']:.2f} L")
        
        # 3. 消除阶段
        elimination_result = self.run_elimination(params, self.time_points, distribution_result)
        print(f"  Elimination completed: total clearance={elimination_result['clearance_rates']['total']:.2f} L/h")
        
        # 返回完整结果
        return {
            'time_points': self.time_points,
            'elimination': elimination_result,
            'models': {
                'absorption': absorption_model,
                'distribution': distribution_model
            }
        }
    
    def calculate_pk_metrics(self, time_points, concentration):
        """Calculate pharmacokinetic parameters"""
        # Cmax和Tmax
        cmax = np.max(concentration)
        tmax = time_points[np.argmax(concentration)]
        
        # AUC (梯形法)
        auc = np.trapezoid(concentration, time_points)
        
        # 终末半衰期
        if len(concentration) >= 10:
            terminal_points = concentration[-10:]
            if np.any(terminal_points > 0):
                log_conc = np.log(terminal_points[terminal_points > 0])
                time_term = time_points[-len(terminal_points):][terminal_points > 0]
                if len(log_conc) > 1:
                    slope, _ = np.polyfit(time_term, log_conc, 1)
                    half_life = -0.693 / slope if slope < 0 else 999
                else:
                    half_life = 999
            else:
                half_life = 999
        else:
            half_life = 999
        
        return {
            'Cmax': cmax,
            'Tmax': tmax,
            'AUC': auc,
            'Half_life': half_life
        }
    
    def plot_results(self, results_dict):
        """Plot results comparison"""
        # Create separate figure for each model combination
        fig, axes = plt.subplots(2, 2, figsize=(15, 12))
        axes = axes.flatten()
        
        # Plot each model combination in its own subplot
        model_combinations = [
            ('first_order', 'minimal_pbpk'),
            ('first_order', 'full_pbpk'),
            ('adam', 'minimal_pbpk'),
            ('adam', 'full_pbpk')
        ]
        
        colors = ['blue', 'red', 'green', 'purple']
        
        for idx, (abs_model, dist_model) in enumerate(model_combinations):
            key = f"{abs_model}_{dist_model}"
            ax = axes[idx]
            
            if key in results_dict:
                result = results_dict[key]
                ax.plot(result['time_points'], result['elimination']['plasma_concentration'], 
                       color=colors[idx], linewidth=2)
                ax.set_xlabel('Time (h)')
                ax.set_ylabel('Plasma Concentration (mg/L)')
                ax.set_title(f'{abs_model.title()} Absorption + {dist_model.replace("_", " ").title()} Distribution')
                ax.grid(True, alpha=0.3)
                
                # Add PK parameters as text
                metrics = self.calculate_pk_metrics(
                    result['time_points'],
                    result['elimination']['plasma_concentration']
                )
                pk_text = f'Cmax: {metrics["Cmax"]:.3f} mg/L\nTmax: {metrics["Tmax"]:.1f} h\nAUC: {metrics["AUC"]:.1f} mg·h/L'
                if metrics['Half_life'] < 100:
                    pk_text += f'\nHalf-life: {metrics["Half_life"]:.1f} h'
                ax.text(0.05, 0.95, pk_text, transform=ax.transAxes, 
                       verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
            else:
                ax.text(0.5, 0.5, 'No Data', ha='center', va='center', fontsize=12)
                ax.set_xlabel('Time (h)')
                ax.set_ylabel('Plasma Concentration (mg/L)')
                ax.set_title(f'{abs_model.title()} Absorption + {dist_model.replace("_", " ").title()} Distribution')
        
        plt.tight_layout()
        plt.savefig('../output/individual_pbpk_results.png', dpi=300, bbox_inches='tight')
        plt.show()


def main():
    """Main function"""
    print("=" * 60)
    print("Individual PBPK Modeling System (Simplified Version)")
    print("=" * 60)
    
    # 创建模拟器
    simulator = SimplePBPKSimulator()
    
    # 定义要运行的模型组合
    model_combinations = [
        ('first_order', 'minimal_pbpk'),
        ('first_order', 'full_pbpk'),
        ('adam', 'minimal_pbpk'),
        ('adam', 'full_pbpk')
    ]
    
    # 存储所有结果
    all_results = {}
    
    # 运行所有模型组合
    for abs_model, dist_model in model_combinations:
        key = f"{abs_model}_{dist_model}"
        
        try:
            result = simulator.run_complete_pbpk(abs_model, dist_model)
            all_results[key] = result
            
            # 打印PK参数
            metrics = simulator.calculate_pk_metrics(
                result['time_points'],
                result['elimination']['plasma_concentration']
            )
            
            print(f"\n  PK parameters:")
            print(f"    Cmax: {metrics['Cmax']:.3f} mg/L")
            print(f"    Tmax: {metrics['Tmax']:.1f} h")
            print(f"    AUC: {metrics['AUC']:.1f} mg·h/L")
            print(f"    Half-life: {metrics['Half_life']:.1f} h" if metrics['Half_life'] < 100 else "    Half-life: N/A")
            
        except Exception as e:
            print(f"  Error: {str(e)}")
            continue
    
    # 生成图表
    if all_results:
        print("\nGenerating plots...")
        simulator.plot_results(all_results)
        print("\nSimulation completed! Plot saved as: individual_pbpk_results.png")
    else:
        print("\nNo models ran successfully!")


if __name__ == "__main__":
    main()